'''
Created on Feb 18, 2013

@author: nzhao
'''
import numpy as np
import scipy.linalg as li

from traits.api import Instance, Float, List, Int

from math_qm.QMobj import QMobject
from math_qm.QMoperator import QMoperator
from math_qm.QMstate import QMState

from phy.solver.eigen import MultiLevel_EigenSystem

class InstantOperation(QMoperator):
    pass

class CoherentEvolve(QMoperator):
    hami = Instance(np.matrix)
    ker = Instance(MultiLevel_EigenSystem)
    time = Float()
    
    def __init__(self, eig, t, is_auto_update=True):
        self.ker = eig
        self.time = t
        if is_auto_update:
            self.calculate_transform_mat()
            
    def calculate_transform_mat(self):
        print self.ker.mat
        print self.time
        self.hami = self.ker.mat
        self.mat = np.matrix( li.expm2(- 1.j * self.hami * self.time) )

class StateEvolution(QMobject):
    initial_state = Instance(QMState)
    final_state   = Instance(QMState)
    operations    = List( QMoperator )
    n_operation   = Int()
    
    def __init__(self, initial_state, operations, is_auto_update=True):
        QMobject.__init__(self,initial_state.nbody, initial_state.dim_list)
        self.initial_state = initial_state
        self.operations     = operations
        self.n_operator    = len(operations)
        if is_auto_update:
            self.calculate_final_state()
    
    def calculate_final_state(self):
        state_mat = self.initial_state.mat
        for op in self.operations:
            state_mat = op.mat.H * state_mat * op.mat
        self.final_state = QMState(self.nbody, self.dim_list, mat= state_mat)
        
        
if __name__ == '__main__':
    from math_qm.vector import vector
    from phy.interaction.SingleSpinInteraction import SingleSpinInteraction
    from phy.interaction.MultiSpinInteraction import MultiSpinInteraction
    from phy.solver.eigen import Eigen_NSpinHalf

    
    applied_b = vector(xyz=[0., 0., 100.])
    
    s1=SingleSpinInteraction(element='C13', spin_type='Bath', x=0.0, y=0.0, z=0.0, b=applied_b)
    state1 = QMState(qubit_state='thermal')
    
    s2=SingleSpinInteraction(element='C13', spin_type='Bath', x=0.0, y=0.0, z=1.54, b=applied_b)
    state2 = QMState(qubit_state='thermal')
    spins = MultiSpinInteraction(spins=[s1,s2])
    
    evo= CoherentEvolve( Eigen_NSpinHalf(spins), 1.0 )
    init_dm = QMState(entries = [state1, state2])
    print 
    
    se = StateEvolution(initial_state=init_dm, operations=[evo])
    print se.final_state.mat
    
    
    
    
    
    
    
    
    
